package com.yangzb.sample.hadoop.mapreduce;

import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.*;
import org.apache.hadoop.io.IOUtils;
import org.apache.hadoop.security.UserGroupInformation;
import org.eclipse.jetty.util.StringUtil;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.net.URI;
import java.util.ArrayList;
import java.util.List;

@Slf4j
public class HadoopUtil extends SecurityUtil {

    public static final String HADOOP_SECURITY_AUTHENTICATION_KEY = "hadoop.security.authentication";
    public static final String HADOOP_USER_NAME = "HADOOP_USER_NAME";
    public static final String HADOOP_HOME_DIR = "hadoop.home.dir";

    /**
     * kerberos 认证
     *
     * @param kerberosKeytabFilePath Keytab文件路径
     * @param kerberosPrincipal      principal （klist -key XX.keytab可以查看）
     * @param krb5Realm              /etc/krb5.autoconfigure 中的 realms
     * @param kdc                    etc/krb5.autoconfigure 中 realms.kdc
     *                               <br>return type: void
     * @Title: kerberosAuthentication
     * @Description:
     * @author yzb yangzb@tydic.com,yzhengbin@gmail.com
     * @date 2017年7月13日 下午5:04:57
     */
    public static boolean kerberosAuthentication(String kerberosKeytabFilePath,
                                                 String kerberosPrincipal, String krb5Realm, String kdc) {
        Configuration hadoopConf = new Configuration();
        return kerberosAuthentication(hadoopConf, kerberosKeytabFilePath, kerberosPrincipal, krb5Realm,
                kdc);
    }
    public static boolean kerberosAuthentication(String kerberosKeytabFilePath,
                                                 String kerberosPrincipal, String krb5Realm, String kdc,String krb5ConfFile) {
        Configuration hadoopConf = new Configuration();
        return kerberosAuthentication(hadoopConf, kerberosKeytabFilePath, kerberosPrincipal, krb5Realm,
                kdc);
    }

    public static boolean kerberosAuthentication(Configuration hadoopConf, String kerberosKeytabFilePath, String kerberosPrincipal, String krb5Realm, String kdc,String krb5ConfFile) {
        hadoopConf.set(HADOOP_SECURITY_AUTHENTICATION_KEY, "kerberos");
        SecurityUtil.setJavaSecurityConf(krb5Realm, kdc,krb5ConfFile);
        UserGroupInformation.setConfiguration(hadoopConf);
        try {
            UserGroupInformation.loginUserFromKeytab(kerberosPrincipal, kerberosKeytabFilePath);
            UserGroupInformation.getLoginUser().checkTGTAndReloginFromKeytab();
            return UserGroupInformation.isLoginKeytabBased();
        } catch (Exception e) {
            String message = String
                    .format("kerberos认证失败,请确定kerberosKeytabFilePath[%s]和kerberosPrincipal[%s]填写正确",
                            kerberosKeytabFilePath, kerberosPrincipal);
            log.error(message, e);
        }
        return false;
    }
    /**
     * kerberos 认证
     *
     * @param kerberosKeytabFilePath Keytab文件路径
     * @param kerberosPrincipal      principal （klist -key XX.keytab可以查看）
     * @param krb5Realm              /etc/krb5.autoconfigure 中的 realms
     * @param kdc                    etc/krb5.autoconfigure 中 realms.kdc
     *                               <br>return type: void
     * @Title: kerberosAuthentication
     * @Description:
     * @author yzb yangzb@tydic.com,yzhengbin@gmail.com
     * @date 2017年7月13日 下午5:04:57
     */
    public static boolean kerberosAuthentication(Configuration hadoopConf, String kerberosKeytabFilePath, String kerberosPrincipal, String krb5Realm, String kdc) {
        hadoopConf.set(HADOOP_SECURITY_AUTHENTICATION_KEY, "kerberos");
        SecurityUtil.setJavaSecurityConf(krb5Realm, kdc);
        UserGroupInformation.setConfiguration(hadoopConf);
        try {
            UserGroupInformation.loginUserFromKeytab(kerberosPrincipal, kerberosKeytabFilePath);
            UserGroupInformation.getLoginUser().checkTGTAndReloginFromKeytab();
            return UserGroupInformation.isLoginKeytabBased();
        } catch (Exception e) {
            String message = String
                    .format("kerberos认证失败,请确定kerberosKeytabFilePath[%s]和kerberosPrincipal[%s]填写正确",
                            kerberosKeytabFilePath, kerberosPrincipal);
            log.error(message, e);
        }
        return false;
    }

    /**
     * @param userName <br>return type: void
     * @Title: setHadoopUser
     * @Description:
     * @author yzb yangzb@tydic.com,yzhengbin@gmail.com
     * @date 2017年7月14日 上午10:47:53
     */
    public static void setHadoopUser(String userName) {
        System.setProperty(HADOOP_USER_NAME, userName);
    }

    /**
     * @param hadoopHomeDir <br>return type: void
     * @Title: setHadoopHomeDir
     * @Description:
     * @author yzb yangzb@tydic.com,yzhengbin@gmail.com
     * @date 2017年7月14日 上午10:48:52
     */
    public static void setHadoopHomeDir(String hadoopHomeDir) {
        System.setProperty(HADOOP_HOME_DIR, hadoopHomeDir);
    }

    /**
     * @param kerberosKeytabFilePath Keytab文件路径
     * @param kerberosPrincipal      principal （klist -key XX.keytab可以查看）
     *                               <br>return type: void
     * @Title: kerberosAuthentication
     * @Description:
     * @author yzb yangzb@tydic.com,yzhengbin@gmail.com
     * @date 2017年7月13日 下午5:08:04
     */
    public static boolean kerberosAuthentication(String kerberosKeytabFilePath,
                                                 String kerberosPrincipal) {
        return kerberosAuthentication(kerberosKeytabFilePath, kerberosPrincipal, null, null);
    }


    /**
     * 获取文件系统
     *
     * @return FileSystem 文件系统
     */
    @SneakyThrows
    public static FileSystem getFileSystem(String hdfsUri) {
        //读取配置文件
        Configuration conf = new Configuration();
        // 文件系统
        FileSystem fs = null;
        if (StringUtil.isBlank(hdfsUri)) {
            try {
                fs = FileSystem.get(conf);
            } catch (IOException e) {
                log.error(e.getMessage(), e);
                log.error("", e);
            }
        } else {
            // 返回指定的文件系统,如果在本地测试，需要使用此种方法获取文件系统
            try {
                URI uri = new URI(hdfsUri.trim());
                fs = FileSystem.get(uri, conf);
            } catch (Exception e) {
                throw e;
            }
        }
        return fs;
    }

    /**
     * 判断目录是否存在
     *
     * @param hdfs     上传服务器
     * @param filePath 目录路径
     * @param create   若不存在是否创建
     */
    @SneakyThrows
    public static boolean existDir(String hdfs, String filePath, boolean create){
        boolean flag = false;
        if (StringUtil.isBlank(filePath)) {
            return flag;
        }
        try {
            Path path = new Path(filePath);
            // FileSystem对象
            FileSystem fs = getFileSystem(hdfs);

            if (create) {
                if (!fs.exists(path)) {
                    fs.mkdirs(path);
                }
            }

            if (fs.isDirectory(path)) {
                flag = true;
            }
        } catch (Exception e) {
            throw e;
        }
        return flag;
    }

    /**
     * 本地文件上传至 HDFS
     *
     * @param srcFile  源文件 路径
     * @param destPath hdfs路径
     */
    @SneakyThrows
    public static void copyFileToHDFS(String hdfsUri, String srcFile, String destPath){
        Configuration hdfsConfig = new Configuration();

        FileInputStream fis = new FileInputStream(new File(srcFile));//读取本地文件
        FileSystem fs = FileSystem.get(URI.create(hdfsUri + destPath), hdfsConfig);
        OutputStream out = null;
        System.out.println(fs.isDirectory(new Path(destPath)));
        if (!fs.isDirectory(new Path(destPath))) {
            out = fs.create(new Path(destPath));
        }
        //copy
        IOUtils.copyBytes(fis, out, 4096, true);
        fs.close();
    }

    /**
     * 获取hdfs中一个目录下的所有子文件名称，比如00000_0
     *
     * @param hdfsUri    hdfs路径
     * @param parentPath 父目录路径
     * @return
     * @throws Exception
     */
    @SneakyThrows
    public static List<String> getHdfsSubFilePath(String hdfsUri, String parentPath){
        Configuration hdfsConfig = new Configuration();
        List<String> subFilePaths = new ArrayList<>();
        try (FileSystem fs = FileSystem.get(URI.create(hdfsUri + parentPath), hdfsConfig)) {
            if (!fs.isDirectory(new Path(parentPath))) {
                throw new Exception("请传一个目录路径！");
            }
            FileStatus[] fileStatuses = fs.listStatus(new Path(parentPath));
            for (FileStatus fileStatus : fileStatuses) {
                if (!fileStatus.isDirectory()) {
                    subFilePaths.add(fileStatus.getPath().getName());
                }
            }
        } catch (Exception e) {
            log.error("获取hdfs中一个目录下的儿子文件名称路径出错", e);
            throw e;
        }
        return subFilePaths;
    }

    @SneakyThrows
    public static List<FileStatus> getHdfsSubFile(String hdfsUri, String parentPath){
        Configuration hdfsConfig = new Configuration();
        List<FileStatus> subFilePaths = new ArrayList<>();
        try (FileSystem fs = FileSystem.get(URI.create(hdfsUri + parentPath), hdfsConfig)) {
            if (!fs.isDirectory(new Path(parentPath))) {
                throw new Exception("请传一个目录路径！");
            }
            FileStatus[] fileStatuses = fs.listStatus(new Path(parentPath));
            for (FileStatus fileStatus : fileStatuses) {
                if (!fileStatus.isDirectory()) {
                    subFilePaths.add(fileStatus);
                }
            }
        } catch (Exception e) {
            log.error("获取hdfs中一个目录下的儿子文件名称路径出错", e);
            throw e;
        }
        return subFilePaths;
    }

    /**
     * 删除文件或者文件目录
     */
    @SneakyThrows
    public static void rmdir(String hdfsUri,String path) {
        try {
            // 返回FileSystem对象
            FileSystem fs = getFileSystem(hdfsUri);
            fs.delete(new Path(path),true);
            fs.close();
        } catch (Exception e) {
            throw e;
        }
    }


}
